// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#ifndef _ST_HPC_PPL_NN_ENGINES_ARM_ARM_DEVICE_H_
#define _ST_HPC_PPL_NN_ENGINES_ARM_ARM_DEVICE_H_

#include "ppl/common/generic_cpu_allocator.h"
#include "ppl/nn/common/device.h"
#include <cstring>

namespace ppl { namespace nn { namespace arm {

class ArmDevice : public Device {
public:
    ArmDevice(uint64_t alignment, ppl::common::isa_t isa) : isa_(isa), allocator_(alignment) {
        *(uint64_t*)(type_.str) = 0;
        type_.str[0] = 'c';
        type_.str[1] = 'p';
        type_.str[2] = 'u';
    }

    void SetISA(ppl::common::isa_t isa) {
        isa_ = isa;
    }
    const ppl::common::isa_t GetISA() const {
        return isa_;
    }

    virtual ppl::common::RetCode AllocTmpBuffer(uint64_t bytes, BufferDesc* buffer) {
        return Realloc(bytes, buffer);
    }

    virtual void FreeTmpBuffer(BufferDesc* buffer) {
        Free(buffer);
    }

    ppl::common::Allocator* GetAllocator() const {
        return &allocator_;
    }

    ppl::common::RetCode Realloc(uint64_t bytes, BufferDesc* buffer) override {
        if (buffer->addr) {
            allocator_.Free(buffer->addr);
        }

        if (bytes == 0) {
            buffer->addr = nullptr;
            return ppl::common::RC_SUCCESS;
        }

        buffer->addr = allocator_.Alloc(bytes);
        if (!buffer->addr) {
            return ppl::common::RC_OUT_OF_MEMORY;
        }

        buffer->desc = bytes;
        return ppl::common::RC_SUCCESS;
    }

    void Free(BufferDesc* buffer) override {
        if (buffer->addr) {
            allocator_.Free(buffer->addr);
            buffer->addr = nullptr;
        }
    }

    ppl::common::RetCode Realloc(const TensorShape& shape, BufferDesc* buffer) override final {
        return Realloc(shape.CalcBytesIncludingPadding(), buffer);
    }

    ppl::common::RetCode CopyFromHost(BufferDesc* dst, const void* src, uint64_t bytes) const override final {
        memcpy(dst->addr, src, bytes);
        return ppl::common::RC_SUCCESS;
    }
    ppl::common::RetCode CopyFromHostAsync(BufferDesc* dst, const void* src, uint64_t bytes) const override final {
        return CopyFromHost(dst, src, bytes);
    }

    ppl::common::RetCode CopyFromHost(BufferDesc* dst, const void* src, const TensorShape& shape) const override final {
        return CopyFromHost(dst, src, shape.CalcBytesIncludingPadding());
    }
    ppl::common::RetCode CopyFromHostAsync(BufferDesc* dst, const void* src,
                                           const TensorShape& shape) const override final {
        return CopyFromHost(dst, src, shape);
    }

    ppl::common::RetCode CopyToHost(void* dst, const BufferDesc& src, uint64_t bytes) const override final {
        memcpy(dst, src.addr, bytes);
        return ppl::common::RC_SUCCESS;
    }
    ppl::common::RetCode CopyToHostAsync(void* dst, const BufferDesc& src, uint64_t bytes) const override final {
        return CopyToHost(dst, src, bytes);
    }

    ppl::common::RetCode CopyToHost(void* dst, const BufferDesc& src, const TensorShape& shape) const override final {
        return CopyToHost(dst, src, shape.CalcBytesIncludingPadding());
    }
    ppl::common::RetCode CopyToHostAsync(void* dst, const BufferDesc& src,
                                         const TensorShape& shape) const override final {
        return CopyToHost(dst, src, shape);
    }

    ppl::common::RetCode Copy(BufferDesc* dst, const BufferDesc& src, uint64_t bytes) const override final {
        memcpy(dst->addr, src.addr, bytes);
        return ppl::common::RC_SUCCESS;
    }

    ppl::common::RetCode Copy(BufferDesc* dst, const BufferDesc& src, const TensorShape& shape) const override final {
        return Copy(dst, src, shape.CalcBytesIncludingPadding());
    }

    ppl::common::RetCode ConvertToHost(void* dst, const TensorShape& dst_desc, const BufferDesc& src,
                                       const TensorShape& src_desc,
                                       const void* src_custom_info = nullptr) override final;
    ppl::common::RetCode ConvertToHostAsync(void* dst, const TensorShape& dst_desc, const BufferDesc& src,
                                            const TensorShape& src_desc,
                                            const void* src_custom_info = nullptr) override final;

    ppl::common::RetCode ConvertFromHost(BufferDesc* dst, const TensorShape& dst_desc, const void* src,
                                         const TensorShape& src_desc,
                                         const void* dst_custom_info = nullptr) override final;
    ppl::common::RetCode ConvertFromHostAsync(BufferDesc* dst, const TensorShape& dst_desc, const void* src,
                                              const TensorShape& src_desc,
                                              const void* dst_custom_info = nullptr) override final;

    ppl::common::RetCode Convert(BufferDesc* dst, const TensorShape& dst_desc, const BufferDesc& src,
                                 const TensorShape& src_desc, const void* dst_custom_info = nullptr,
                                 const void* src_custom_info = nullptr) override final;

    ppl::common::RetCode Synchronize() override final {
        return ppl::common::RC_SUCCESS;
    }

    const Type& GetType() const override final {
        return type_;
    }

    ppl::common::RetCode Configure(uint32_t, ...) override {
        return ppl::common::RC_UNSUPPORTED;
    }

private:
    bool MayUseISA(uint32_t flag) const {
        return !!(isa_ & flag);
    }

private:
    Type type_;
    ppl::common::isa_t isa_;
    mutable ppl::common::GenericCpuAllocator allocator_;
};

}}} // namespace ppl::nn::arm

#endif
